import numpy as np
from Tools import metrics_result
from Tools import readbunchobj
from sklearn.svm import SVC
from joblib import dump
from sklearn.preprocessing import StandardScaler

# 导入训练集
trainpath = "../train_word_bag/tf_idf_space.dat"
train_set = readbunchobj(trainpath)

# 导入测试集
testpath = "../test_word_bag/test_tf_idf_space.dat"
test_set = readbunchobj(testpath)

# 输出单词矩阵的类型
print("标准化/归一化前的矩阵形状:")
print(np.shape(train_set.tdm))
print(np.shape(test_set.tdm))

# 对数据进行标准化或归一化
scaler = StandardScaler(with_mean=False)
train_set.tdm = scaler.fit_transform(train_set.tdm)
test_set.tdm = scaler.transform(test_set.tdm)

# 支持向量机
print("开始训练支持向量机模型...")
SVM_clf = SVC(kernel='linear', C=1.0, random_state=0).fit(train_set.tdm, train_set.label)
print("支持向量机模型训练完成。")
dump(SVM_clf, 'SVM_model.joblib')

# 预测分类结果
print("USE: 支持向量机模型 预测分类结果")
SVM_predicted = SVM_clf.predict(test_set.tdm)
SVM_total = len(SVM_predicted)
print("结束")

# 性能评估
print("\n支持向量机模型：")
metrics_result(test_set.label, SVM_predicted, SVM_total)
